#pragma once
#include "common.h"
#include <cstring>

// 单层数组
template <int BITS>
class TCMalloc_PageMap1
{
private:
    static const int LENGTH = 1 << BITS;
    void **array_;

public:
    typedef uintptr_t Number;

    // 显式构造函数
    explicit TCMalloc_PageMap1()
    {
        size_t size = sizeof(void *) << BITS;
        size_t alignSize = SizeClass::_RoundUp(size, 1 << PAGE_SHIFT);
        array_ = (void **)SystemAlloc(alignSize >> PAGE_SHIFT);
        memset(array_, 0, sizeof(void *) << BITS);
    }

    // 返回键KEY的当前值。如果尚未设置或k超出范围，则返回NULL
    void *get(Number k) const
    {
        if ((k >> BITS) > 0)
        {
            return NULL;
        }
        return array_[k];
    }

    // 要求"k"在范围"[0,2^BITS-1]"内
    // 要求"k"之前已经被确保存在
    //
    // 为键"k"设置值"v"
    void set(Number k, void *v)
    {
        array_[k] = v;
    }
};

// 两基数树
template <int BITS>
class TCMalloc_PageMap2
{
private:
    // 在根节点放置32个条目，在每个叶子节点放置(2^BITS)/32个条目
    static const int ROOT_BITS = 5;
    static const int ROOT_LENGTH = 1 << ROOT_BITS;

    static const int LEAF_BITS = BITS - ROOT_BITS;
    static const int LEAF_LENGTH = 1 << LEAF_BITS;

    // 叶子节点结构
    struct Leaf
    {
        void *values[LEAF_LENGTH];
    };

    Leaf *root_[ROOT_LENGTH]; // 指向32个子节点的指针

public:
    typedef uintptr_t Number;

    // 显式构造函数
    explicit TCMalloc_PageMap2()
    {
        memset(root_, 0, sizeof(root_));

        PreallocateMoreMemory();
    }

    // 获取键k对应的值
    void *get(Number k) const
    {
        const Number i1 = k >> LEAF_BITS;
        const Number i2 = k & (LEAF_LENGTH - 1);
        if ((k >> BITS) > 0 || root_[i1] == NULL)
        {
            return NULL;
        }
        return root_[i1]->values[i2];
    }

    // 设置键k对应的值为v
    void set(Number k, void *v)
    {
        const Number i1 = k >> LEAF_BITS;
        const Number i2 = k & (LEAF_LENGTH - 1);
        assert(i1 < ROOT_LENGTH);
        root_[i1]->values[i2] = v;
    }

    // 确保从start开始的n个连续键已经准备好可以使用
    bool Ensure(Number start, size_t n)
    {
        for (Number key = start; key <= start + n - 1;)
        {
            const Number i1 = key >> LEAF_BITS;

            // 检查是否溢出
            if (i1 >= ROOT_LENGTH)
                return false;

            // 必要时创建二级节点
            if (root_[i1] == NULL)
            {
                static ObjectPool<Leaf> leafPool;
                Leaf *leaf = (Leaf *)leafPool.New();

                memset(leaf, 0, sizeof(*leaf));
                root_[i1] = leaf;
            }

            // 将键推进到下一个叶子节点覆盖的范围
            key = ((key >> LEAF_BITS) + 1) << LEAF_BITS;
        }
        return true;
    }

    // 预分配更多内存
    void PreallocateMoreMemory()
    {
        // 分配足够的内存来跟踪所有可能的页面
        Ensure(0, 1 << BITS);
    }
};